import torch
import numpy as np
import os
from torch.utils.data import Dataset, DataLoader
import glob
import os
import numpy as np
import nibabel as nib
import torch
from torch.utils.data import Dataset, DataLoader
import random
from sklearn.model_selection import train_test_split
import albumentations as A

class Brats2D(Dataset):
    def __init__(self,caselist,train=True,datadir="/mnt/data1/MedicalDataset/BraTS18_2D/") -> None:
        self.datadir=datadir
        self.cases=caselist
        self.slicelist=[]
        self.train=train
        for case in self.cases:
            case_slices=os.listdir(self.datadir+case)
            for case_slice in case_slices:
                self.slicelist.append(os.path.join(self.datadir+case,case_slice))
        
    def __getitem__(self, index):
        """
        Return corresponding flair, t1, t1ce, t2, seg_gt
        """
        flair,t1,t1ce,t2,seg_gt=np.load(self.slicelist[index])
        if self.train==True:
            temp=np.concatenate([flair,t1,t1ce,t2,seg_gt])
            temp=np.transpose(temp,(1,2,0))
            transform = A.Compose([
                A.RandomCrop(width=224,height=224),
                A.Flip(p=0.5)
                ])
            temp=transform(image=temp)['image']
            temp=np.transpose(temp,(2,0,1))
            flair,t1,t1ce,t2,seg_gt=temp[0],temp[1],temp[2],temp[3],temp[4]
        flair,t1,t1ce,t2=flair/flair.max()*2-1,t1/t1.max()*2-1,t1ce/t1ce.max()*2-1,t2/t2.max()*2-1
        #return  torch.tensor(flair.copy(), dtype=torch.float).unsqueeze(0),torch.tensor(t1.copy(), dtype=torch.float).unsqueeze(0),torch.zeros_like(torch.tensor(t1ce.copy(), dtype=torch.float).unsqueeze(0)),torch.tensor(t2.copy(), dtype=torch.float).unsqueeze(0),torch.tensor(seg_gt.copy(), dtype=torch.float).unsqueeze(0)
        wt_volume = seg_gt > 0  # 坏死和无增强的肿瘤区域：1、增强区域（活跃部分）：4、周边水肿区域：2
        tc_volume = np.logical_or(seg_gt == 4, seg_gt == 1)
        et_volume = (seg_gt == 4)
        seg_volume = [[tc_volume, wt_volume, et_volume]]
        seg_volume = np.concatenate(seg_volume, axis=0).astype("float32")
        return  torch.tensor(flair.copy(), dtype=torch.float).unsqueeze(0),torch.tensor(t1.copy(), dtype=torch.float).unsqueeze(0),torch.tensor(t1ce.copy(), dtype=torch.float).unsqueeze(0),torch.tensor(t2.copy(), dtype=torch.float).unsqueeze(0),torch.tensor(seg_volume.copy(), dtype=torch.float)
    def __len__(self):
        return len(self.slicelist)

def split_dataset(data_root="/mnt/data1/MedicalDataset/BraTS18_2D", nfold=5, seed=42, select=0):
    n_patients =os.listdir(data_root)
    train_patients_list, val_patients_list=train_test_split(n_patients,test_size=0.2, random_state=42)
    return train_patients_list, val_patients_list

